import os
import cv2
import json
import numpy as np

from glob import glob
from tqdm import tqdm
from PIL import Image

from .dataset import Dataset
from .video import Video


class DepthTrackVideo(Video):
    def __init__(self, name, root, video_dir, init_rect, img_names,
                 gt_rect, load_img=False):
        super(DepthTrackVideo, self).__init__(name, root, video_dir,
                                              init_rect, img_names, gt_rect, None, load_img)
        self.gt_traj = [[0] if np.isnan(bbox[0]) else bbox
                        for bbox in self.gt_traj]
        if not load_img:
            img_name = self.img_names[0]
            img = np.array(Image.open(img_name), np.uint8)
            self.width = img.shape[1]
            self.height = img.shape[0]
        self.confidence = {}

    def load_tracker(self, path, tracker_names=None, store=True):
        """
        Args:
            path(str): path to result
            tracker_name(list): name of tracker
        """
        if not tracker_names:
            tracker_names = [x.split('/')[-1] for x in glob(path)
                             if os.path.isdir(x)]
        if isinstance(tracker_names, str):
            tracker_names = [tracker_names]
        for name in tracker_names:
            traj_file = os.path.join(path, name, self.name + '.txt')
            with open(traj_file, 'r') as f:
                traj = [list(map(float, x.strip().split(',')))
                        for x in f.readlines()]
            if store:
                self.pred_trajs[name] = traj
            confidence_file = os.path.join(path, name, self.name + '_confidence.txt')
            with open(confidence_file, 'r') as f:
                score = [float(x.strip()) for x in f.readlines()[1:]]
                score.insert(0, float('nan'))
            if store:
                self.confidence[name] = score
        return traj, score


class DepthTrackDataset(Dataset):
    def __init__(self, name, dataset_root, load_img=False):
        super(DepthTrackDataset, self).__init__(name, dataset_root)
        with open(os.path.join(dataset_root, name + '.json'), 'r') as f:
            meta_data = json.load(f)

        # load videos
        pbar = tqdm(meta_data.keys(), desc='loading ' + name, ncols=100)
        self.videos = {}
        for video in pbar:
            pbar.set_postfix_str(video)
            self.videos[video] = DepthTrackVideo(video,
                                                 dataset_root,
                                                 meta_data[video]['video_dir'],
                                                 meta_data[video]['init_rect'],
                                                 meta_data[video]['img_names'],
                                                 meta_data[video]['gt_rect'])
